import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * <b>description</b>： Java进阶、性能优化实战案例100讲，微信公众号:程序员恰恰，源码 & 文档 & 技术支持，请加个人微信号：it3306 <br>
 * <b>time</b>：2025/4/25 13:27 <br>
 * <b>author</b>：ready lao_gao
 */

// 继承JDK里面的饿ThreadPoolExecutor
public class MyThreadPoolExecutor extends ThreadPoolExecutor {

    private ThreadLocalContext threadLocalContext;

    public MyThreadPoolExecutor(ThreadLocalContext context, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
        this.threadLocalContext = context;
    }

    public MyThreadPoolExecutor(ThreadLocalContext context, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
        this.threadLocalContext = context;
    }

    //重写了 execute 方法，在提交任务时，将原始任务包装为 RunnableWrap对象。
    // RunnableWrap 的作用是捕获主线程的上下文，并在任务执行时将其传递到线程池中的线程。
    @Override
    public void execute(Runnable command) {
        super.execute(this.new RunnableWrap(command));
    }

    //重写了 submit 方法，在提交 Callable 任务时，将原始任务包装为 CallableWrap 对象。
    // CallableWrap 的作用与 RunnableWrap 类似，也是为了捕获和传递上下文。
    @Override
    public <T> Future<T> submit(Callable<T> task) {
        return super.submit(this.new CallableWrap<>(task));
    }

    /**
     * 线程本地变量上下文接口，用于解决线程池中共享外部线程ThreadLocal数据的问题
     *
     * @param <T>
     */
    public interface ThreadLocalContext<T> {
        /**
         * 从主线程中获取当前的上下文对象。
         *
         * @return
         */
        T getContext();

        /**
         * 将上下文对象设置到线程池中的线程。
         *
         * @param context {@link #getContext()} 返回的对象
         */
        void setContext(T context);
    }

    /**
     * 捕获主线程的上下文，并在任务执行时将其传递到线程池中的线程。
     *
     * @param <V>
     */
    private class CallableWrap<V> implements Callable<V> {

        private Callable<V> target;

        private Object context;

        public CallableWrap(Callable<V> target) {
            this.target = target;
            //在构造函数中，通过 threadLocalContext.getContext() 捕获主线程的上下文数据
            this.context = MyThreadPoolExecutor.this.threadLocalContext.getContext();
        }

        @Override
        public V call() throws Exception {
            //在 run 方法中，通过 threadLocalContext.setContext(this.context)
            // 将上下文数据设置到线程池中的线程，然后执行原始任务。
            MyThreadPoolExecutor.this.threadLocalContext.setContext(this.context);
            return this.target.call();
        }
    }


    private class RunnableWrap implements Runnable {

        private Runnable target;

        private Object context;

        public RunnableWrap(Runnable target) {
            this.target = target;
            this.context = MyThreadPoolExecutor.this.threadLocalContext.getContext();
        }

        @Override
        public void run() {
            //在 call 方法中，先设置上下文数据，然后调用原始任务的 call 方法。
            MyThreadPoolExecutor.this.threadLocalContext.setContext(this.context);
            this.target.run();
        }
    }
}